package com.snda.dw.lidejia.partitionfilegenerator;

import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.io.HiveKey;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.mapred.FileInputFormat;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.mapred.lib.InputSampler;
import org.apache.hadoop.mapred.lib.TotalOrderPartitioner;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;

/**
 * Hello world!
 *
 */
public class PartitionerFileGenerator extends Configured implements Tool {

    public static void main(String[] args) throws Exception {
        int code = 0;
        if (args.length > 0) {
//            code = ToolRunner.run(new PartitionFileGenerator(), args[0].split(" "));
            code = ToolRunner.run(new PartitionerFileGenerator(), args);
        } else {
            System.err.println("Missing arguments");
            code = -1;
        }
        System.exit(code);

    }

    public int run(String[] args) throws Exception {
        final String filePath = args[0];
        final String[] partitionColumns = new String[args.length - 1];
        System.arraycopy(args, 1, partitionColumns, 0, partitionColumns.length);
        //
        JobConf conf = new JobConf(getConf());
        //
        conf.setInputFormat(HiveKeyFileInputFormat.class);
        FileInputFormat.addInputPath(conf, new Path(filePath));
        conf.setMapOutputKeyClass(HiveKey.class);
        conf.setMapOutputValueClass(NullWritable.class);
        // set partition columns
        StringBuilder partitions = new StringBuilder();
        for (String partitionColumn : partitionColumns) {
            partitions.append(partitionColumn);
            partitions.append(";");
        }
        partitions.deleteCharAt(partitions.length() - 1);
        conf.set("totalsort.partitioner.columns", partitions.toString());
        // set partition number
        conf.setNumReduceTasks(5);
        // set partition path
        TotalOrderPartitioner.setPartitionFile(conf, new Path("/tmp/_totalsort_partitions"));
        // init sampler
        InputSampler.Sampler<HiveKey, NullWritable> sampler = new InputSampler.RandomSampler<HiveKey, NullWritable>(0.1, 10000, 50);
        // get partitions
        InputSampler.writePartitionFile(conf, sampler);
        //
        return 0;
    }
}
